package com.smart.translator.util;

import android.util.Log;
import java.io.ByteArrayOutputStream;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.ShortBuffer;

/**
 * 音频重采样器 - 将8kHz PCM转换为16kHz PCM
 * 使用线性插值算法进行重采样
 */
public class AudioResampler {
    private static final String TAG = "AudioResampler";
    
    // 采样率常量
    public static final int SAMPLE_RATE_8K = 8000;
    public static final int SAMPLE_RATE_16K = 16000;
    
    // 重采样比例
    private static final double RESAMPLE_RATIO = (double) SAMPLE_RATE_16K / SAMPLE_RATE_8K; // 2.0
    
    // 缓冲区
    private ByteArrayOutputStream outputBuffer;
    private byte[] tempBuffer;
    
    public AudioResampler() {
        outputBuffer = new ByteArrayOutputStream();
        tempBuffer = new byte[4096]; // 4KB临时缓冲区
    }
    
    /**
     * 将8kHz PCM数据转换为16kHz PCM数据
     * @param inputData 输入的8kHz PCM数据
     * @return 转换后的16kHz PCM数据
     */
    public byte[] resample8kTo16k(byte[] inputData) {
        if (inputData == null || inputData.length == 0) {
            return new byte[0];
        }
        
        try {
            // 将字节数组转换为short数组（16位PCM）
            ShortBuffer inputBuffer = ByteBuffer.wrap(inputData)
                    .order(ByteOrder.LITTLE_ENDIAN)
                    .asShortBuffer();
            
            short[] inputShorts = new short[inputBuffer.remaining()];
            inputBuffer.get(inputShorts);
            
            // 计算输出样本数
            int inputSamples = inputShorts.length;
            int outputSamples = (int) Math.ceil(inputSamples * RESAMPLE_RATIO);
            
            // 执行重采样
            short[] outputShorts = linearResample(inputShorts, outputSamples);
            
            // 转换回字节数组
            ByteBuffer outputBuffer = ByteBuffer.allocate(outputShorts.length * 2)
                    .order(ByteOrder.LITTLE_ENDIAN);
            outputBuffer.asShortBuffer().put(outputShorts);
            
            // Log.d(TAG, "Resampled " + inputSamples + " samples to " + outputSamples + 
            //            " samples (" + inputData.length + " bytes -> " + outputBuffer.array().length + " bytes)");
            
            return outputBuffer.array();
            
        } catch (Exception e) {
            Log.e(TAG, "Error during resampling", e);
            return inputData; // 出错时返回原始数据
        }
    }
    
    /**
     * 线性插值重采样
     * @param input 输入样本
     * @param outputSamples 输出样本数
     * @return 重采样后的样本
     */
    private short[] linearResample(short[] input, int outputSamples) {
        short[] output = new short[outputSamples];
        
        for (int i = 0; i < outputSamples; i++) {
            // 计算输入样本的索引（浮点数）
            double inputIndex = i / RESAMPLE_RATIO;
            
            // 获取相邻的两个输入样本
            int index1 = (int) Math.floor(inputIndex);
            int index2 = Math.min(index1 + 1, input.length - 1);
            
            // 计算插值权重
            double weight = inputIndex - index1;
            
            // 线性插值
            if (index1 < input.length) {
                short sample1 = input[index1];
                short sample2 = (index2 < input.length) ? input[index2] : sample1;
                
                output[i] = (short) (sample1 * (1 - weight) + sample2 * weight);
            } else {
                output[i] = 0;
            }
        }
        
        return output;
    }
    
    /**
     * 批量处理音频数据
     * @param inputData 输入的8kHz PCM数据
     * @param isLast 是否为最后一块数据
     * @return 转换后的16kHz PCM数据
     */
    public byte[] processAudioData(byte[] inputData, boolean isLast) {
        if (inputData == null || inputData.length == 0) {
            return new byte[0];
        }
        
        // 执行重采样
        byte[] resampledData = resample8kTo16k(inputData);
        
        // 添加到输出缓冲区
        try {
            outputBuffer.write(resampledData);
        } catch (Exception e) {
            Log.e(TAG, "Error writing to output buffer", e);
        }
        
        return resampledData;
    }
    
    /**
     * 获取缓冲区中的所有数据
     * @return 缓冲区中的所有16kHz PCM数据
     */
    public byte[] getBufferedData() {
        byte[] data = outputBuffer.toByteArray();
        outputBuffer.reset();
        return data;
    }
    
    /**
     * 清空缓冲区
     */
    public void clearBuffer() {
        outputBuffer.reset();
    }
    
    /**
     * 获取缓冲区大小
     * @return 缓冲区中的字节数
     */
    public int getBufferSize() {
        return outputBuffer.size();
    }
    
    /**
     * 计算重采样后的数据大小
     * @param originalSize 原始8kHz数据大小（字节）
     * @return 重采样后的16kHz数据大小（字节）
     */
    public static int calculateResampledSize(int originalSize) {
        return (int) Math.ceil(originalSize * RESAMPLE_RATIO);
    }
    
    /**
     * 检查数据大小是否为有效的PCM数据
     * @param dataSize 数据大小（字节）
     * @return 是否为有效的PCM数据大小
     */
    public static boolean isValidPcmSize(int dataSize) {
        return dataSize > 0 && dataSize % 2 == 0; // 16位PCM，每个样本2字节
    }
} 